import java.util.Arrays;
import java.util.PriorityQueue;

public class Solution {
    /**
     * 在 O(n)的时间复杂度之内，在array数组中寻找最小的k个元素
     * @param array
     */
    public static Object[] smallestK(int[] array,int k) {
        //判断array是否为null或者为空数组
        if(array == null || array.length == 0) {
            throw new RuntimeException("数组为空");
        }

        if(k <= 0 || k > array.length) {
            throw new RuntimeException("k不合法");
        }

        Object[] result;
        //建立一个大根堆
        PriorityQueue<Integer> pq = new PriorityQueue<>(new IntCmp());
        for (int i = 0; i < k; i++) {
            pq.offer(array[i]);
        }

        //遍历剩下的n-k个元素，如果有元素小于堆顶元素，则弹出堆顶元素，然后将此元素加入堆
        for (int i = k; i < array.length; i++) {
            if(array[i] < pq.peek()) {
                pq.poll();
                pq.offer(array[i]);
            }
        }

        //算法走到这里，优先队列中已经存储了数组中最小的k个元素
        //将结果复制到result数组中
        result = pq.toArray();
        Arrays.sort(result);
        return result;
    }

    public static Object[] biggestK(int[] array,int k) {
        //判断array是否为null或者为空数组
        if(array == null || array.length == 0) {
            throw new RuntimeException("数组为空");
        }

        if(k <= 0 || k > array.length) {
            throw new RuntimeException("k不合法");
        }

        Object[] result;
        //建立一个小根堆
        PriorityQueue<Integer> pq = new PriorityQueue<>();
        for (int i = 0; i < k; i++) {
            pq.offer(array[i]);
        }

        //遍历剩下的n - k个元素，如果有元素大于堆顶元素，则弹出堆顶元素，然后将此元素加入堆
        for (int i = k; i < array.length; i++) {
            if(array[i] > pq.peek()) {
                pq.poll();
                pq.offer(array[i]);
            }
        }

        result = pq.toArray();
        Arrays.sort(result);
        return result;
    }

    public static void main(String[] args) {
        int[] array = new int[]{7,6,5,4,3,2,1};
        int k = 2;
        System.out.println(Arrays.toString(smallestK(array, k)));
        System.out.println(Arrays.toString(biggestK(array,k)));
    }
}
